/*
 * Copyright 2020 Red Hat, Inc. and/or its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.kie.memorycompiler;

import java.net.URL;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

import org.drools.util.PortablePath;
import org.kie.memorycompiler.resources.MemoryResourceReader;
import org.kie.memorycompiler.resources.MemoryResourceStore;

import static org.kie.memorycompiler.JavaConfiguration.findJavaVersion;

public class KieMemoryCompiler {

    private KieMemoryCompiler() { }

    /**
     * Compile the given sources <b>without</b> adding compiled classes to the given <code>ClassLoader</code>
     * <b>classNameSourceMap</b>' key must be the <b>FQDN</b> of the class to compile
     *
     * @param classNameSourceMap
     * @param classLoader
     * @return
     */
    public static Map<String, Class<?>> compile(Map<String, String> classNameSourceMap, ClassLoader classLoader) {
        return compile(classNameSourceMap, classLoader, null);
    }

    /**
     * Compile the given sources <b>without</b> adding compiled classes to the given <code>ClassLoader</code>
     * <b>classNameSourceMap</b>' key must be the <b>FQDN</b> of the class to compile.
     * Additional compiler settings can be provided using JavaCompilerSettings
     *
     * @param classNameSourceMap
     * @param classLoader
     * @param compilerSettings
     * @return
     */
    public static Map<String, Class<?>> compile(Map<String, String> classNameSourceMap, ClassLoader classLoader, JavaCompilerSettings compilerSettings) {
        Map<String, byte[]> byteCode = compileNoLoad(classNameSourceMap, classLoader, compilerSettings);

        MemoryCompilerClassLoader kieMemoryCompilerClassLoader = new MemoryCompilerClassLoader(classLoader);

        Map<String, Class<?>> toReturn = new HashMap<>();
        for (Map.Entry<String, byte[]> entry : byteCode.entrySet()) {
            kieMemoryCompilerClassLoader.addCode( entry.getKey(), entry.getValue() );
            try {
                toReturn.put(entry.getKey(), kieMemoryCompilerClassLoader.loadClass(entry.getKey()));
            } catch (ClassNotFoundException e) {
                throw new KieMemoryCompilerException(e.getMessage(), e);
            }
        }
        return toReturn;
    }

    /**
     * Compile the given sources and returns the generated byte codes
     *
     * @param classNameSourceMap
     * @param classLoader
     * @return
     */
    public static Map<String, byte[]> compileNoLoad(Map<String, String> classNameSourceMap, ClassLoader classLoader) {
        return compileNoLoad(classNameSourceMap, classLoader, null, JavaConfiguration.CompilerType.NATIVE);
    }

    /**
     * Compile the given sources and returns the generated byte codes
     *
     * @param classNameSourceMap
     * @param classLoader
     * @return
     */
    public static Map<String, byte[]> compileNoLoad(Map<String, String> classNameSourceMap, ClassLoader classLoader, JavaConfiguration.CompilerType compilerType) {
        return compileNoLoad(classNameSourceMap, classLoader, null, compilerType);
    }

    /**
     * Compile the given sources and returns the generated byte codes.
     * Additional compiler settings can be provided using JavaCompilerSettings
     *
     * @param classNameSourceMap
     * @param classLoader
     * @param compilerSettings
     * @return
     */
    public static Map<String, byte[]> compileNoLoad(Map<String, String> classNameSourceMap, ClassLoader classLoader, JavaCompilerSettings compilerSettings) {
        return compileNoLoad(classNameSourceMap, classLoader, compilerSettings, JavaConfiguration.CompilerType.NATIVE);
    }

    /**
     * Compile the given sources and returns the generated byte codes.
     * Additional compiler settings can be provided using JavaCompilerSettings
     *
     * @param classNameSourceMap
     * @param classLoader
     * @param compilerSettings
     * @return
     */
    public static Map<String, byte[]> compileNoLoad(Map<String, String> classNameSourceMap, ClassLoader classLoader, JavaCompilerSettings compilerSettings, JavaConfiguration.CompilerType compilerType) {
        MemoryResourceReader reader = new MemoryResourceReader();
        MemoryResourceStore store = new MemoryResourceStore();
        String[] classNames = new String[classNameSourceMap.size()];

        int i = 0;
        for (Map.Entry<String, String> entry : classNameSourceMap.entrySet()) {
            classNames[i] = toJavaSource( entry.getKey() );
            reader.add( classNames[i], entry.getValue().getBytes());
            i++;
        }
        JavaConfiguration javaConfiguration = new JavaConfiguration();
        javaConfiguration.setCompiler(compilerType);
        javaConfiguration.setJavaLanguageLevel(findJavaVersion());
        JavaCompiler compiler = JavaCompilerFactory.loadCompiler(javaConfiguration);
        CompilationResult res = compilerSettings == null ?
                compiler.compile( classNames, reader, store, classLoader) :
                compiler.compile( classNames, reader, store, classLoader, compilerSettings);

        if (res.getErrors().length > 0) {
            throw new KieMemoryCompilerException(Arrays.toString( res.getErrors() ));
        }

        Map<String, byte[]> toReturn = new HashMap<>();
        for (Map.Entry<PortablePath, byte[]> entry : store.getResources().entrySet()) {
            toReturn.put(toClassName( entry.getKey().asString() ), entry.getValue());
        }

        return toReturn;
    }

    private static String toJavaSource( String s ) {
        return s.replace( '.', '/' ) + ".java";
    }

    private static String toClassName( String s ) {
        if (s.endsWith(".class")) {
            s = s.substring(0, s.length()-6);
        }
        return s.replace( '/', '.' );
    }

    public static class MemoryCompilerClassLoader extends ClassLoader {

        private Map<String, byte[]> customCompiledCode = new HashMap<>();
        private Map<String, URL> customAddedResource = new HashMap<>();

        public MemoryCompilerClassLoader(ClassLoader parent) {
            super(parent);
        }

        public void addCode(String name, byte[] bytes) {
            customCompiledCode.put(name, bytes);
        }
        public void addResource(String name, URL url) {
            customAddedResource.put(name, url);
        }


        public byte[] getCode(String name) {
            return customCompiledCode.get(name);
        }

        public void addCodeIfAbsent(String name, byte[] bytes) {
            customCompiledCode.putIfAbsent(name, bytes);
        }

        @Override
        protected Class<?> findClass(String name) throws ClassNotFoundException {
            byte[] byteCode = customCompiledCode.get(name);
            if (byteCode == null) {
                return super.findClass(name);
            }
            return defineClass(name, byteCode, 0, byteCode.length);
        }

        @Override
        protected URL findResource(String name) {
            URL url = customAddedResource.get(name);
            return url != null ? url :  super.findResource(name);
        }
    }
}
